# AS Strong meta file generation
# as_strong_meta: source meta file from AudioSet Strong
import os
import yaml
import pandas as pd
from itertools import chain

# Load common labels
common_labels = "./DESED_AS_common_labels.yaml"
with open(common_labels, "r") as f:
    common_labels = yaml.safe_load(f)

id_labels = list(chain.from_iterable(common_labels.values()))

# Load as strong meta
as_root = "/home/shaonian/Datasets/AudioSet_strong/data/train/"
as_strong_meta = "/home/shaonian/Datasets/AudioSet_strong/meta/source/audioset_train_strong.tsv"
as_strong_df = pd.read_csv(as_strong_meta, sep="\t")
out_domain_df = pd.DataFrame({}, columns=["filename"])
in_domain_df = pd.DataFrame({}, columns=["filename"])
invalid_mask = as_strong_df.iloc[:, 3].isin(id_labels)
invalid_filenames = as_strong_df.loc[invalid_mask, "segment_id"].unique()
valid_filenames = list(set(as_strong_df.loc[:, "segment_id"].unique()) - set(invalid_filenames))
print("valid: {}, invalid: {}".format(len(valid_filenames), len(invalid_filenames)))
valid_filenames = [ x + ".wav" for x in valid_filenames ]
invalid_filenames = [ x + ".wav" for x in invalid_filenames ]
# check data existing
valid_file_list = []
for f in valid_filenames:
    abs_path = os.path.join(as_root, f)
    if os.path.exists(abs_path):
        valid_file_list.append(f)

invalid_file_list = []
for f in invalid_filenames:
    abs_path = os.path.join(as_root, f)
    if os.path.exists(abs_path):
        invalid_file_list.append(f)

out_domain_df["filename"] = valid_file_list
out_domain_df.to_csv("./DESED_AS_OD_meta.tsv", index=False, sep="\t")

in_domain_df["filename"] = invalid_file_list
in_domain_df.to_csv("./DESED_AS_ID_meta.tsv", index=False, sep="\t")